// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// Functions for comparing Nebula data structures

#pragma once

#include <cstdint>
#include <iosfwd>

#include <turbo/base/macros.h>

namespace nebula {

    class Array;

    class DataType;

    class Tensor;

    class SparseTensor;

    struct Scalar;

    static constexpr double kDefaultAbsoluteTolerance = 1E-5;

    /// A container of options for equality comparisons
    class EqualOptions {
    public:
        /// Whether or not NaNs are considered equal.
        bool nans_equal() const { return nans_equal_; }

        /// Return a new EqualOptions object with the "nans_equal" property changed.
        EqualOptions nans_equal(bool v) const {
            auto res = EqualOptions(*this);
            res.nans_equal_ = v;
            return res;
        }

        /// Whether or not zeros with differing signs are considered equal.
        bool signed_zeros_equal() const { return signed_zeros_equal_; }

        /// Return a new EqualOptions object with the "signed_zeros_equal" property changed.
        EqualOptions signed_zeros_equal(bool v) const {
            auto res = EqualOptions(*this);
            res.signed_zeros_equal_ = v;
            return res;
        }

        /// The absolute tolerance for approximate comparisons of floating-point values.
        double atol() const { return atol_; }

        /// Return a new EqualOptions object with the "atol" property changed.
        EqualOptions atol(double v) const {
            auto res = EqualOptions(*this);
            res.atol_ = v;
            return res;
        }

        /// The ostream to which a diff will be formatted if arrays disagree.
        /// If this is null (the default) no diff will be formatted.
        std::ostream *diff_sink() const { return diff_sink_; }

        /// Return a new EqualOptions object with the "diff_sink" property changed.
        /// This option will be ignored if diff formatting of the types of compared arrays is
        /// not supported.
        EqualOptions diff_sink(std::ostream *diff_sink) const {
            auto res = EqualOptions(*this);
            res.diff_sink_ = diff_sink;
            return res;
        }

        static EqualOptions defaults() { return {}; }

    protected:
        double atol_ = kDefaultAbsoluteTolerance;
        bool nans_equal_ = false;
        bool signed_zeros_equal_ = true;

        std::ostream *diff_sink_ = nullptr;
    };

    /// Returns true if the arrays are exactly equal
    TURBO_EXPORT bool array_equals(const Array &left, const Array &right,
                                  const EqualOptions & = EqualOptions::defaults());

    /// Returns true if the arrays are approximately equal. For non-floating point
    /// types, this is equivalent to array_equals(left, right)
    TURBO_EXPORT bool array_approx_equals(const Array &left, const Array &right,
                                        const EqualOptions & = EqualOptions::defaults());

    /// Returns true if indicated equal-length segment of arrays are exactly equal
    TURBO_EXPORT bool array_range_equals(const Array &left, const Array &right,
                                       int64_t start_idx, int64_t end_idx,
                                       int64_t other_start_idx,
                                       const EqualOptions & = EqualOptions::defaults());

    /// Returns true if indicated equal-length segment of arrays are approximately equal
    TURBO_EXPORT bool array_range_approx_equals(const Array &left, const Array &right,
                                             int64_t start_idx, int64_t end_idx,
                                             int64_t other_start_idx,
                                             const EqualOptions & = EqualOptions::defaults());

    TURBO_EXPORT bool tensor_equals(const Tensor &left, const Tensor &right,
                                   const EqualOptions & = EqualOptions::defaults());

    /// EXPERIMENTAL: Returns true if the given sparse tensors are exactly equal
    TURBO_EXPORT bool sparse_tensor_equals(const SparseTensor &left, const SparseTensor &right,
                                         const EqualOptions & = EqualOptions::defaults());

    /// Returns true if the type metadata are exactly equal
    /// \param[in] left a DataType
    /// \param[in] right a DataType
    /// \param[in] check_metadata whether to compare KeyValueMetadata for child
    /// fields
    TURBO_EXPORT bool type_equals(const DataType &left, const DataType &right,
                                 bool check_metadata = true);

    /// Returns true if scalars are equal
    /// \param[in] left a Scalar
    /// \param[in] right a Scalar
    /// \param[in] options comparison options
    TURBO_EXPORT bool scalar_equals(const Scalar &left, const Scalar &right,
                                   const EqualOptions &options = EqualOptions::defaults());

    /// Returns true if scalars are approximately equal
    /// \param[in] left a Scalar
    /// \param[in] right a Scalar
    /// \param[in] options comparison options
    TURBO_EXPORT bool scalar_approx_equals(
            const Scalar &left, const Scalar &right,
            const EqualOptions &options = EqualOptions::defaults());

}  // namespace nebula
